import os
import faiss
import pandas as pd
from flask import Flask, request, redirect, render_template, send_from_directory
from werkzeug.utils import secure_filename
from sklearn.feature_extraction.text import TfidfVectorizer

# Flask 配置
app = Flask(__name__)

# 设置文件上传路径
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'csv'}
INDEX_FOLDER = 'faiss_indexes'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['INDEX_FOLDER'] = INDEX_FOLDER

# 确保上传和索引文件夹存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(INDEX_FOLDER, exist_ok=True)


# 允许的文件扩展名
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# 上传并处理 CSV 文件
@app.route('/', methods=['GET'])
def upload_form():
    return render_template('index.html')


@app.route('/upload_csv', methods=['POST'])
def upload_csv():
    # 检查文件是否存在
    if 'file' not in request.files:
        return "No file part"

    file = request.files['file']

    # 如果用户没有选择文件
    if file.filename == '':
        return "No selected file"

    # 检查文件是否为 CSV
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)

        # 处理 CSV 文件并生成 FAISS 索引
        create_faiss_index(file_path, filename)

        return f"File uploaded and FAISS index created successfully: {filename}"

    return "Invalid file type, please upload a CSV file."


# 创建 FAISS 索引
def create_faiss_index(csv_file, filename):
    # 读取 CSV 文件
    df = pd.read_csv(csv_file)

    # 假设 CSV 包含两个列：车牌号和品牌
    if 'plateNo' not in df.columns or 'carBrand' not in df.columns:
        raise ValueError("CSV file must contain 'plateNo' and 'carBrand' columns.")

    # 组合车牌号和品牌信息用于检索
    data = df['plateNo'] + " " + df['carBrand']

    # 向量化
    vectorizer = TfidfVectorizer()
    vectors = vectorizer.fit_transform(data).toarray()

    # 创建 FAISS 索引
    dimension = vectors.shape[1]
    faiss_index = faiss.IndexFlatL2(dimension)
    faiss_index.add(vectors)

    # 保存向量化模型
    index_filename = filename.rsplit('.', 1)[0] + ".index"
    index_path = os.path.join(app.config['INDEX_FOLDER'], index_filename)

    faiss.write_index(faiss_index, index_path)

    return index_path


# 运行 Flask 应用
if __name__ == "__main__":
    app.run(debug=True)
